from pathlib import Path
from typing import Union
import warnings
from torch.utils.data import Dataset


class MyDataset(Dataset):

    def __init__(self, p: Union[Path, str]):
        self.path = Path(p)
        self.label_dir = []
        for entry in self.path.glob('*'):
            if entry.is_dir():
                self.label_dir.append(entry)
                # for label in entry.glob('*'):
                #     if label.suffix.lower() == '.wav':
                #         print(label)
        self.chilidren = {}
        self.fd = None
        self.has_children = False
        self.keys = set()

    def __len__(self):
        return len(self.label_dir)

    def __getitem__(self, idx):
        label_dir = self.label_dir[idx]
        label_name = label_dir.stem
        files = list(label_dir.glob('*'))
        local_wav = [str(file) for file in files if file.suffix.lower() == '.wav']
        label_data = [str(file) for file in files if file.suffix.lower() == '.xlsx']
        label_source_dir = [file for file in files if file.is_dir()][0]
        label_wav = []
        for label in label_source_dir.glob('*'):
            if label.suffix.lower() == '.wav':
                label_wav.append(str(label))
        return {'name': label_name, 'source': local_wav,
                'label': label_wav, 'info': label_data}

    # def __setitem__(self, key: str, value: str):
    #     if self.has_children:
    #         raise RuntimeError("This writer points out a directory")
    #     if key in self.keys:
    #         warnings.warn(f"Duplicated: {key}")
    #
    #     if self.fd is None:
    #         self.path.parent.mkdir(parents=True, exist_ok=True)
    #         self.fd = self.path.open("w", encoding="utf-8")
    #
    #     self.keys.add(key)
    #     self.fd.write(f"{key} {value}\n")
    #     self.fd.flush()
    #
    # def __exit__(self, exc_type, exc_val, exc_tb):
    #     self.close()
    #
    # def close(self):
    #     if self.has_children:
    #         prev_child = None
    #         for child in self.chilidren.values():
    #             child.close()
    #             if prev_child is not None and prev_child.keys != child.keys:
    #                 warnings.warn(
    #                     f"Ids are mismatching between "
    #                     f"{prev_child.path} and {child.path}"
    #                 )
    #             prev_child = child
    #
    #     elif self.fd is not None:
    #         self.fd.close()
